# File: fedentgate/core/server.py
"""Federated Learning Server Implementation"""

import torch
from .attention_aggregation import CoAttentionAggregator

class FLServer:
    def __init__(self, global_model, num_clients):
        """
        Initialize federated learning server
        Args:
            global_model: Initial global model
            num_clients: Total number of clients
        """
        self.global_model = global_model
        self.aggregator = CoAttentionAggregator(input_dim=3)
        self.client_states = {i: {'entropy': 0.0, 'gating': 0} for i in range(num_clients)}
    
    def update_client_state(self, client_id, entropy, gating_state):
        """Update stored client state information"""
        self.client_states[client_id]['entropy'] = entropy
        self.client_states[client_id]['gating'] = gating_state
    
    def aggregate_updates(self, client_updates):
        """
        Aggregate client updates using attention weights
        Args:
            client_updates: Dict {client_id: (gradients, noise_feature)}
        Returns:
            avg_update: Aggregated model update
        """
        # Build input vectors for attention computation
        input_vectors = []
        for cid, update in client_updates.items():
            state = self.client_states[cid]
            vec = self.aggregator.build_input_vector(
                state['entropy'], 
                state['gating'], 
                update[1]  # noise_feature
            )
            input_vectors.append(vec)
        
        # Compute attention weights
        attn_weights = self.aggregator(input_vectors)
        
        # Apply weighted aggregation
        weighted_update = None
        for i, (cid, (grads, _)) in enumerate(client_updates.items()):
            weight = attn_weights[i]
            if weighted_update is None:
                weighted_update = [g * weight for g in grads]
            else:
                for j, g in enumerate(grads):
                    weighted_update[j] += g * weight
        
        return weighted_update
